import tensorflow.keras as keras
import json
from preprocess import SEQUENCE_LENGTH, MAPPING_PATH
from train import SAVE_MODEL_PATH
import numpy as np
import music21 as m21


class MelodyGenerator:
    
    def __init__(self, model_path=SAVE_MODEL_PATH):
        
        self.model_path = model_path
        self.model = keras.models.load_model(model_path)
        
        with open(MAPPING_PATH, 'r') as fp:
            self._mappings = json.load(fp)
            
        self._start_symbols = ['/'] * SEQUENCE_LENGTH
    
    
    def _sample_with_temperature(self, probabilities, temperature):
        """
        how creative the model will be
        
        ``temperature`` -> infinity: randomly pick the probability
        
        ``temperature`` -> 0: pick the highest probability
        
        ``temperature`` = 1: normal sampling, use the distribution received from the network
        
        """
        predictions = np.log(probabilities) / temperature
        probabilities = np.exp(predictions) / np.sum(np.exp(predictions))
        
        choices = range(len(probabilities))
        index = np.random.choice(choices, p=probabilities)
        
        return index
            
    
    def generate_melody(self, seed, num_steps, max_sequence_length, temperature):
        """Parameters:
        
        `seed`: melody sequences
        
        `num_steps`: number of steps the network will generate to output
        
        `max_sequence_length`: limited part of seed pass to the network, determin how many steps to pass to the network. 64
        
        `temperature`: how createive the model will be. 
        
        """
        # create seed with start symbols
        seed = seed.split()
        melody = seed
        seed = self._start_symbols + seed
        
        # map seed to integers
        seed = [self._mappings[symbol] for symbol in seed]
        
        for _ in range(num_steps):
            
            # limit the seed to the max sequence length
            seed = seed[-max_sequence_length:]
            
            # one-hot encode the seed
            # create a new dimension because keras expects a batch of samples
            # (1, max_sequence_length, num of symbols in the vocabulary)
            onehot_seed = keras.utils.to_categorical(seed, num_classes=len(self._mappings))
            onehot_seed = onehot_seed[np.newaxis, ...]
            
            # make a prediction
            probabilities = self.model.predict(onehot_seed)[0]  # a batch of probabilities, just want the first
            output_int = self._sample_with_temperature(probabilities, temperature)
            
            # update seed
            seed.append(output_int)
            
            # map int to our encoding
            output_symbol = [k for k, v in self._mappings.items() if v == output_int][0]
            
            # check whether we're at the end of a melody
            if output_symbol == '/':
                break
            
            # update the melody
            melody.append(output_symbol)
        
        return melody

        
    def save_melody(self, melody, step_duration=0.25, format='midi', file_name='mel_shanxi.mid'):
        """保存旋律到midi文件
        """
        # create a music21 stream
        stream = m21.stream.Stream()
        
        # parse all the symbols in the melody and create note/rest objects
        # 60 _ _ _ r _ 62 _
        start_symbol = None
        step_counter = 1  # track steps for an event
        
        for i, symbol in enumerate(melody):
            
            # handle case in which we have a note/rest
            if symbol != '_' or i + 1 == len(melody):
                
                # ensure we're dealing with note/rest beyond the first one
                if start_symbol is not None:
                    
                    quarter_length_duration = step_duration * step_counter
                    
                    # handle rest
                    if start_symbol == 'r':
                        m21_event = m21.note.Rest(quarterLength=quarter_length_duration)
                        
                    # handle notes
                    else:
                        m21_event = m21.note.Note(int(start_symbol), quarterLength=quarter_length_duration)
                        
                    stream.append(m21_event)
                    
                    # reset the step counter
                    step_counter = 1
                    
                # update the start symbol
                start_symbol = symbol
                                
            # handle case in which we have a prolongation sign(_)
            else:
                step_counter += 1
                
        # write the m21 stream to a midi file
        stream.write(format, file_name)
            

if __name__ == "__main__":
    mg = MelodyGenerator()
    seed = '74 _ _ 74 74 _ 79 _ 69 _ _'
    seed2 = '67 _ _ _ _ _ 65 _ 64 _ 62 _ 60 _ _ _'
    melody = mg.generate_melody(seed, 500, SEQUENCE_LENGTH, 0.7)
    print(melody)
    mg.save_melody(melody)
    